import pyrallis

from src.configs.train_config import TrainConfig
from src.training.trainer import TEXTure


@pyrallis.wrap()
def main(cfg: TrainConfig):
    trainer = TEXTure(cfg)
    
    if cfg.log.mode == 'overlap_latent':
        trainer.paint_overlap_latent()
        exit()
    elif cfg.log.mode == 'overlap_img_x0':
        trainer.paint_overlap_imgx0()
        exit()
    elif cfg.log.mode == 'overlap_img_xt':
        trainer.paint_overlap_imgxt()
        exit()

    # if cfg.log.eval_only:
    #     trainer.full_eval()
    # else:
    #     trainer.paint()


if __name__ == '__main__':
    main()
